import os
import argparse
from datetime import datetime

def create_test_files(directory, count, prefix="test"):
    """创建指定数量的测试文件"""
    for i in range(1, count+1):
        file_path = os.path.join(directory, f"{prefix}_{i}.txt")
        with open(file_path, 'w') as f:
            f.write(f"测试文件 #{i}\n创建时间: {datetime.now()}")

def create_test_environment(root_dir="test_data"):
    """
    创建针对压缩功能优化的测试环境
    参数:
    - root_dir: 测试根目录名称
    """
    # 优化后的测试场景配置
    scenarios = {
        'threshold_edge': {'files': 5, 'dirs': 1},  # 刚好达到阈值
        'above_threshold': {'files': 8, 'dirs': 2},  # 明显超过阈值
        'below_threshold': {'files': 3, 'dirs': 1},  # 低于阈值
        'deep_nested': {'files': 2, 'dirs': 5, 'depth': 4},  # 深层嵌套
        'large_files': {'files': 15, 'dirs': 0},  # 大量文件无子目录
        'mixed_content': {'files': 4, 'dirs': 4}  # 混合内容
    }
    
    # 创建根目录
    project_dir = os.path.dirname(os.path.abspath(__file__))
    test_root = os.path.join(project_dir, root_dir)
    os.makedirs(test_root, exist_ok=True)
    print(f"测试根目录: {test_root}")

    # 创建各种测试场景
    for scenario, config in scenarios.items():
        scenario_dir = os.path.join(test_root, f"scenario_{scenario}")
        os.makedirs(scenario_dir, exist_ok=True)
        print(f"\n创建测试场景: {scenario}")
        
        # 创建基础文件
        create_test_files(scenario_dir, config['files'], f"{scenario}_file")
        
        # 创建子目录结构
        if 'depth' in config:  # 处理深层嵌套
            current_dir = scenario_dir
            for depth in range(1, config['depth']+1):
                new_dir = os.path.join(current_dir, f"level{depth}")
                os.makedirs(new_dir, exist_ok=True)
                create_test_files(new_dir, 1)
                current_dir = new_dir
        else:  # 普通子目录
            for i in range(1, config['dirs']+1):
                sub_dir = os.path.join(scenario_dir, f"subdir_{i}")
                os.makedirs(sub_dir, exist_ok=True)
                create_test_files(sub_dir, 2)  # 每个子目录2个文件

    print("\n优化测试环境创建完成！")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='创建针对压缩优化的测试目录结构')
    parser.add_argument('-d', '--directory', type=str, default="test_data",
                       help='测试根目录名称 (默认: "test_data")')
    args = parser.parse_args()
    
    create_test_environment(args.directory)
